"""
使用matplotlib绘制K线图和双均线图
功能：弹出窗口显示两幅图，上图显示K线，下图显示两条均线
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from mplfinance.original_flavor import candlestick_ohlc
from tqsdk import TqApi, TqAuth, TqKq
from datetime import datetime

# MyTT模块中的函数定义

# 这些是从MyTT库中提取的常用技术指标函数
def REF(Series, N=1):
    """向前引用N周期的数据"""
    return pd.Series(Series).shift(N)

def EMA(Series, N=10):
    """指数移动平均线"""
    return pd.Series(Series).ewm(span=N, adjust=False).mean()

def calculate_signals(klines):
    """
    计算交易信号并生成均线数据
    :param klines: K线数据
    :return: 包含均线和K线数据的DataFrame
    """
    # 转换为numpy数组以提高计算效率
    close = np.array(klines.close)
    # print(klines.close)
    low = np.array(klines.low)
    open_price = np.array(klines.open)
    high = np.array(klines.high)
    
    # 计算Q值
    Q = (3 * close + low + open_price + high) / 6

    #print(Q)
    
    # 计算trading_line
    terms = [
        26 * Q,
        25 * REF(Q, 1),
        24 * REF(Q, 2),
        23 * REF(Q, 3),
        22 * REF(Q, 4),
        21 * REF(Q, 5),
        20 * REF(Q, 6),
        19 * REF(Q, 7),
        18 * REF(Q, 8),
        17 * REF(Q, 9),
        16 * REF(Q, 10),
        15 * REF(Q, 11),
        14 * REF(Q, 12),
        13 * REF(Q, 13),
        12 * REF(Q, 14),
        11 * REF(Q, 15),
        10 * REF(Q, 16),
        9 * REF(Q, 17),
        8 * REF(Q, 18),
        7 * REF(Q, 19),
        6 * REF(Q, 20),
        5 * REF(Q, 21),
        4 * REF(Q, 22),
        3 * REF(Q, 23),
        2 * REF(Q, 24),
        REF(Q, 25)
    ]
    
    trading_line = sum(terms) / 351
    short_line = EMA(trading_line, 7)

    # print(trading_line)
    
    # 创建结果DataFrame
    result = pd.DataFrame({
        'datetime': klines.datetime,
        'open': open_price,
        'high': high,
        'low': low,
        'close': close,
        'trading_line': trading_line,
        'short_line': short_line
    })
    
    # 转换时间戳为日期时间
    try:
        result['date'] = pd.to_datetime(result['datetime'], unit='s')
    except Exception:
        result['date'] = pd.to_datetime(result['datetime'])
    result['date_ordinal'] = result['date'].map(mdates.date2num)
    
    #print(result)

    return result

def create_visualization(symbol, period_name, result):
    """
    创建K线图和双均线图
    :param symbol: 合约代码
    :param period_name: 周期名称
    :param result: 计算结果DataFrame
    """
    # 创建一个包含两个子图的图表
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10), sharex=True)
    fig.suptitle(f"{symbol} {period_name} K线图和双均线图", fontsize=16)
    
    # 准备OHLC数据
    ohlc = result[['date_ordinal', 'open', 'high', 'low', 'close']]
    
    # 绘制K线图
    candlestick_ohlc(ax1, ohlc.values, width=0.6, colorup='g', colordown='r', alpha=0.8)
    ax1.set_title('K线图')
    ax1.set_ylabel('价格')
    ax1.grid(True)
    
    # 绘制均线图
    ax2.plot(result['date'], result['trading_line'], label='Trading Line', color='blue', linewidth=1.5)
    ax2.plot(result['date'], result['short_line'], label='Short Line', color='red', linewidth=1.5)
    ax2.set_title('双均线图')
    ax2.set_ylabel('价格')
    ax2.legend()
    ax2.grid(True)
    
    # 格式化x轴日期
    ax2.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    ax2.xaxis.set_major_locator(mdates.AutoDateLocator())
    plt.setp(ax2.xaxis.get_majorticklabels(), rotation=45)
    
    # 调整布局
    plt.tight_layout()
    plt.subplots_adjust(top=0.92)
    
    # 显示图表
    plt.show()

# 主函数
def main():
    try:
        # 创建API连接
        api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba"))
        
        # 设置默认合约和周期
        symbol = "SHFE.ao2601"  # 上期所黄金
        period = 86400  # 日线
        period_name = "日线"
        
        print(f"正在获取 {symbol} 的{period_name}数据...")
        klines = api.get_kline_serial(symbol, period, data_length=500)
        
        # 计算信号和均线数据
        result = calculate_signals(klines)
        
        # 创建可视化图表
        create_visualization(symbol, period_name, result)
        
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        if 'api' in locals():
            api.close()
            print("API连接已关闭")

if __name__ == "__main__":
    main()